import functools

from coin.base.param_util import to_list
from coin.base.datetime_util import to_datetime
from coin.exchange.base.tick import FixedTickPrice
from coin.strategy.hk_hard_arb.dump_unhedge_mm import FeatureRow
from coin.strategy.mm.calculator import MovingSum


def _gen_for_same_quotes(config):
  amt_det_notional = config['price_det_notional']
  book_askt_func_0 = (
      lambda book: book.get_notional_asks_by_amt(amt_det_notional).entries[-1].price)
  book_askt_func_1 = book_askt_func_0
  book_bidt_func_0 = (
      lambda book: book.get_notional_bids_by_amt(amt_det_notional).entries[-1].price)
  book_bidt_func_1 = book_bidt_func_0
  return book_askt_func_0, book_askt_func_1, book_bidt_func_0, book_bidt_func_1


def _gen_for_different_quotes2(config):
  assert 'price_det_notional' not in config
  assert 'ref_qty_det_notional' in config or 'ref_price_det_notional' in config
  assert 'trade_qty_det_notional' in config or 'trade_price_det_notional' in config

  if 'ref_qty_det_notional' in config:
    qty_det_notional = to_list(config['ref_qty_det_notional'])
    book_askt_func_0 = [
        functools.partial(lambda notional, book: book.get_ask_price_by_qty(notional), det)
        for det in qty_det_notional
    ]
    book_bidt_func_0 = [
        functools.partial(lambda notional, book: book.get_bid_price_by_qty(notional), det)
        for det in qty_det_notional
    ]
  elif 'ref_price_det_notional' in config:
    amt_det_notional = to_list(config['ref_price_det_notional'])
    book_askt_func_0 = [
        functools.partial(lambda notional, book: book.get_ask_price_by_amount(notional), det)
        for det in amt_det_notional
    ]
    book_bidt_func_0 = [
        functools.partial(lambda notional, book: book.get_bid_price_by_amount(notional), det)
        for det in amt_det_notional
    ]
  else:
    raise ValueError('invalid ref')

  if 'trade_qty_det_notional' in config:
    qty_det_notional = config['trade_qty_det_notional']
    book_askt_func_1 = (lambda book: book.get_ask_price_by_qty(qty_det_notional))
    book_bidt_func_1 = (lambda book: book.get_bid_price_by_qty(qty_det_notional))
  elif 'trade_price_det_notional' in config:
    amt_det_notional = config['trade_price_det_notional']
    book_askt_func_1 = (lambda book: book.get_ask_price_by_amount(amt_det_notional))
    book_bidt_func_1 = (lambda book: book.get_bid_price_by_amount(amt_det_notional))
  else:
    raise ValueError('invalid ref')
  return book_askt_func_0, book_askt_func_1, book_bidt_func_0, book_bidt_func_1


def _gen_for_different_quotes(config):
  assert 'price_det_notional' not in config
  assert 'ref_qty_det_notional' in config or 'ref_price_det_notional' in config
  assert 'trade_qty_det_notional' in config or 'trade_price_det_notional' in config

  if 'ref_qty_det_notional' in config:
    qty_det_notional = to_list(config['ref_qty_det_notional'])
    book_askt_func_0 = [
        functools.partial(
            lambda notional, book: book.get_notional_asks_by_qty(notional).entries[-1].price, det)
        for det in qty_det_notional
    ]
    book_bidt_func_0 = [
        functools.partial(
            lambda notional, book: book.get_notional_bids_by_qty(notional).entries[-1].price, det)
        for det in qty_det_notional
    ]
  elif 'ref_price_det_notional' in config:
    amt_det_notional = to_list(config['ref_price_det_notional'])
    book_askt_func_0 = [
        functools.partial(
            lambda notional, book: book.get_notional_asks_by_amt(notional).entries[-1].price, det)
        for det in amt_det_notional
    ]
    book_bidt_func_0 = [
        functools.partial(
            lambda notional, book: book.get_notional_bids_by_amt(notional).entries[-1].price, det)
        for det in amt_det_notional
    ]
  else:
    raise ValueError('invalid ref')

  if 'trade_qty_det_notional' in config:
    qty_det_notional = config['trade_qty_det_notional']
    book_askt_func_1 = (
        lambda book: book.get_notional_asks_by_qty(qty_det_notional).entries[-1].price)
    book_bidt_func_1 = (
        lambda book: book.get_notional_bids_by_qty(qty_det_notional).entries[-1].price)
  elif 'trade_price_det_notional' in config:
    amt_det_notional = config['trade_price_det_notional']
    book_askt_func_1 = (
        lambda book: book.get_notional_asks_by_amt(amt_det_notional).entries[-1].price)
    book_bidt_func_1 = (
        lambda book: book.get_notional_bids_by_amt(amt_det_notional).entries[-1].price)
  else:
    raise ValueError('invalid ref')
  return book_askt_func_0, book_askt_func_1, book_bidt_func_0, book_bidt_func_1


def _gen_true_book_funcs(config, use_fastfeed=False):
  if 'price_det_notional' in config:
    return _gen_for_same_quotes(config)
  else:
    if use_fastfeed:
      return _gen_for_different_quotes2(config)
    else:
      return _gen_for_different_quotes(config)


def _calculate_spread(askp_1, bidp_1):
  return (askp_1 - bidp_1) / (bidp_1 + askp_1) * 2


class IndexCalculator(object):
  def __init__(self, product_askt_funcs, product_bidt_funcs, weights=None):
    assert isinstance(product_askt_funcs, list)
    assert isinstance(product_bidt_funcs, list)
    assert len(product_askt_funcs) == len(product_bidt_funcs)
    if isinstance(weights, list):
      assert len(weights) == len(product_bidt_funcs)
    self._num_products = len(product_askt_funcs)
    self._product_askt_funcs = product_askt_funcs
    self._product_bidt_funcs = product_bidt_funcs
    self._weights = \
      weights or [1 / self._num_products for i in range(self._num_products)]
    self._weighted_cal_func = lambda xs, ys: sum(x * y for x, y in zip(xs, ys))
    self._last = None
    self._ready = False

  def update(self, books):
    assert len(books) == self._num_products
    self._last = books
    self._ready = True

  def get_best_price(self):
    assert self._ready
    ask0_list = [book.ask0().price for book in self._last]
    bid0_list = [book.bid0().price for book in self._last]
    ask0_price = self._weighted_cal_func(ask0_list, self._weights)
    bid0_price = self._weighted_cal_func(bid0_list, self._weights)
    return ask0_price, bid0_price

  def get_true_price(self):
    assert self._ready
    askt_list = \
      [func(book) for book, func in zip(self._last, self._product_askt_funcs)]
    bidt_list = \
      [func(book) for book, func in zip(self._last, self._product_bidt_funcs)]
    askt_price = self._weighted_cal_func(askt_list, self._weights)
    bidt_price = self._weighted_cal_func(bidt_list, self._weights)
    return askt_price, bidt_price


class IndexPricer(object):
  def __init__(self, config, tick, use_fastfeed=False):
    basis_ma_window = (int(config['basis_ma_window_min'] * 60e9))
    ref_book_askt_funcs, trade_book_askt_func, ref_book_bidt_funcs, trade_book_bidt_func = _gen_true_book_funcs(
        config, use_fastfeed)
    price_multiplier_1 = config.get('ccy_0_1', 1.)
    ref_weights = config.get('ref_weights', None)

    self._index = IndexCalculator(ref_book_askt_funcs, ref_book_bidt_funcs, ref_weights)
    self._use_basis = basis_ma_window > 0
    self._basis_ma_window = basis_ma_window
    self._basis_ma_required = 10
    self._tick = FixedTickPrice(tick) if isinstance(tick, float) else tick
    self._trade_book_askt_func = trade_book_askt_func
    self._trade_book_bidt_func = trade_book_bidt_func

    self._ready = False
    self._last = None
    self._price_mult_1 = price_multiplier_1 or 1.

    if self._use_basis:
      self._ask_basis_ma = MovingSum(self._basis_ma_window)
      self._bid_basis_ma = MovingSum(self._basis_ma_window)
      self._mid_basis_ma = MovingSum(self._basis_ma_window)

  @property
  def ready(self):
    return self._ready

  @property
  def last_value(self):
    return self._last

  def reset(self):
    self._ready = False
    self._last = None

  def update(self, ref_books, trade_book):
    assert isinstance(ref_books, list)
    func = lambda book: book.timestamp
    book_ts = max(trade_book, max(ref_books, key=func), key=func).timestamp
    min_ts = min(trade_book, min(ref_books, key=func), key=func).timestamp

    if book_ts - min_ts > (10 * 10**9):
      self.reset()
      return

    all_books = ref_books + [trade_book]
    for book in all_books:
      if not (book.has_bid() and book.has_ask()):
        self.reset()
        return

    # index
    self._index.update(ref_books)
    ask0_index, bid0_index = [
        price * self._price_mult_1 if price is not None else price
        for price in self._index.get_best_price()
    ]
    askt_index, bidt_index = [
        price * self._price_mult_1 if price is not None else price
        for price in self._index.get_true_price()
    ]
    assert askt_index >= ask0_index >= bid0_index >= bidt_index, (askt_index,
                                                                  ask0_index,
                                                                  bid0_index,
                                                                  bidt_index)

    # trade_product
    ask0_trade = trade_book.ask0().price
    bid0_trade = trade_book.bid0().price
    askt_trade = self._trade_book_askt_func(trade_book)
    bidt_trade = self._trade_book_bidt_func(trade_book)
    assert askt_trade >= ask0_trade >= bid0_trade >= bidt_trade, (askt_trade,
                                                                  ask0_trade,
                                                                  bid0_trade,
                                                                  bidt_trade)

    mid0_index = (ask0_index + bid0_index) * 0.5
    mid0_trade = (ask0_trade + bid0_trade) * 0.5

    if self._use_basis:
      askt_diff = askt_index - askt_trade
      bidt_diff = bidt_index - bidt_trade
      midt_diff = mid0_index - mid0_trade

      self._ask_basis_ma.update(book_ts, askt_diff)
      self._bid_basis_ma.update(book_ts, bidt_diff)
      self._mid_basis_ma.update(book_ts, midt_diff)

      if len(self._ask_basis_ma) < self._basis_ma_required or \
         len(self._bid_basis_ma) < self._basis_ma_required:
        self.reset()
        return

      self._last = FeatureRow(time=book_ts,
                              time_hr=str(to_datetime(book_ts)),
                              ask0_0=ask0_index,
                              askt_0=askt_index,
                              bid0_0=bid0_index,
                              bidt_0=bidt_index,
                              ask0_1=ask0_trade,
                              askt_1=askt_trade,
                              bid0_1=bid0_trade,
                              bidt_1=bidt_trade,
                              ask_basis_ma=self._ask_basis_ma.average,
                              bid_basis_ma=self._bid_basis_ma.average)
    else:
      self._last = FeatureRow(time=book_ts,
                              time_hr=str(to_datetime(book_ts)),
                              ask0_0=ask0_index,
                              askt_0=askt_index,
                              bid0_0=bid0_index,
                              bidt_0=bidt_index,
                              ask0_1=ask0_trade,
                              askt_1=askt_trade,
                              bid0_1=bid0_trade,
                              bidt_1=bidt_trade,
                              ask_basis_ma=0.,
                              bid_basis_ma=0.)

    self._ready = True

  def adjust_price(self, askp_1, bidp_1, *, bbo_protection=False):
    askp_1 = self._tick.ceil(askp_1)
    bidp_1 = self._tick.floor(bidp_1)

    ask0_minus_one_1 = self._tick.get_prev_price(self._last.ask0_1)
    bid0_plus_one_1 = self._tick.get_next_price(self._last.bid0_1)

    if bbo_protection:
      askp_1 = max(askp_1, self._last.ask0_1, bid0_plus_one_1)
      bidp_1 = min(bidp_1, ask0_minus_one_1, self._last.bid0_1)
    else:
      askp_1 = max(askp_1, ask0_minus_one_1, bid0_plus_one_1)
      bidp_1 = min(bidp_1, ask0_minus_one_1, bid0_plus_one_1)

    # TODO(inkyu): Deprecate return value
    return askp_1, bidp_1

  def calculate_price(self, sell_edge, buy_edge, *, bbo_protection=False):
    assert self.ready
    assert -0.5 < sell_edge < 0.5
    assert -0.5 < buy_edge < 0.5

    askt_0_proj_1 = self._last.askt_0
    bidt_0_proj_1 = self._last.bidt_0

    if self._use_basis:
      askt_0_proj_1 -= float(self._ask_basis_ma.average)
      bidt_0_proj_1 -= float(self._bid_basis_ma.average)

    askp_1 = askt_0_proj_1 * (1. + sell_edge)
    bidp_1 = bidt_0_proj_1 * (1. - buy_edge)

    return self.adjust_price(askp_1, bidp_1, bbo_protection=bbo_protection)

  def calculate_price_for_arb(self, sell_edge, buy_edge, *, bbo_protection=False):
    assert self.ready
    assert -0.5 < sell_edge < 0.5
    assert -0.5 < buy_edge < 0.5

    askt_1_proj_0 = self._last.askt_1
    bidt_1_proj_0 = self._last.bidt_1

    if self._use_basis:
      askt_1_proj_0 -= float(self._mid_basis_ma.average)
      bidt_1_proj_0 -= float(self._mid_basis_ma.average)

    askp_0 = askt_1_proj_0 * (1. + sell_edge)
    bidp_0 = bidt_1_proj_0 * (1. - buy_edge)

    return self.adjust_price(askp_0, bidp_0, bbo_protection=bbo_protection)

  def calculate_price_by_spread(self, spread, *, bbo_protection=False):
    assert self.ready
    assert 0 < spread < 0.5

    midp_proj_1 = (self._last.askt_0 + self._last.bidt_0) / 2.0

    if self._use_basis:
      midp_proj_1 -= float((self._ask_basis_ma.average + self._bid_basis_ma.average) / 2.0)

    # 0.1bp for buffer
    askp_1 = midp_proj_1 * (1. + spread / 2.0 - 0.00001)
    bidp_1 = midp_proj_1 * (1. - spread / 2.0 + 0.00001)

    cal_spread = _calculate_spread(askp_1, bidp_1)

    assert cal_spread <= spread, (cal_spread, spread)

    return self.adjust_price(askp_1, bidp_1, bbo_protection=bbo_protection)
